import os
import re
import shutil
import argparse
from concurrent.futures import ThreadPoolExecutor

def restructure_directory(directory_path):
    print(f"Processing directory: {directory_path}")
    # Create 'img' subdirectory
    img_dir = os.path.join(directory_path, 'img')
    os.makedirs(img_dir, exist_ok=True)

    # Get all files in the directory
    files = os.listdir(directory_path)

    # Move image files to 'img' subdirectory
    image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]
    for image_file in image_files:
        src = os.path.join(directory_path, image_file)
        dst = os.path.join(img_dir, image_file)
        shutil.move(src, dst)

    # Update markdown file
    markdown_files = [f for f in files if f.lower().endswith('.md')]
    for md_file in markdown_files:
        md_path = os.path.join(directory_path, md_file)
        update_markdown_file(md_path)

def update_markdown_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()

    # Update image paths in markdown content
    updated_content = re.sub(r'!\[(.*?)\]\((?!http)(.*?)\)', r'![\1](img/\2)', content)

    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(updated_content)

def process_input_directory(input_dir):
    # Get all subdirectories in the input directory
    subdirectories = [os.path.join(input_dir, d) for d in os.listdir(input_dir)
                      if os.path.isdir(os.path.join(input_dir, d))]

    # Use ThreadPoolExecutor to process directories in parallel
    with ThreadPoolExecutor() as executor:
        executor.map(restructure_directory, subdirectories)

def main():
    parser = argparse.ArgumentParser(description="Restructure directories and update markdown files.")
    parser.add_argument("input_dir", help="Path to the input directory containing subdirectories to process")
    args = parser.parse_args()

    input_dir = args.input_dir
    if not os.path.isdir(input_dir):
        print(f"错误 Error: {input_dir}  不是一个可行的文件夹  is not a valid directory")
        return

    process_input_directory(input_dir)
    print("Processing complete.")

if __name__ == "__main__":
    main()


# python3 restructure_directories.py /path/to/your/input_dir

    #
#
# python3 -m document2llm.md.marker.restructure_directories     </path/to/your/input_dir>

# # python3 -m document2llm.md.marker.restructure_directories  
#
